import numpy as np


class MyDataSet:
    def __init__(self, years, prices, batch_size, shuffle=True):
        self.years = years
        self.prices = prices
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self):  # 循环时候自动触发，返回一个具有__next__方法的对象。用来循环DataSet
        return MyDataLoader(self)

    def __len__(self):
        assert len(self.years) == len(self.prices)
        return len(self.years)


class MyDataLoader:
    def __init__(self, dataset):
        self.dataSet = dataset
        self.indexs = [i for i in range(len(self.dataSet))]  # 用到了上面的__len__
        if self.dataSet.shuffle:
            np.random.shuffle(self.indexs)
        self.cursor = 0

    def __next__(self):  # 循环完一个位置到下一个位置。
        # 如果到头了，就不用到下一个了，抛出异常退出
        if self.cursor >= len(self.dataSet):
            raise StopIteration

        # 不然就取出batch_size大小的数据
        index = self.indexs[self.cursor: self.cursor + self.dataSet.batch_size]

        # 取出index数组里对应的下标
        years_data = self.dataSet.years[index]
        prices_data = self.dataSet.prices[index]

        self.cursor += self.dataSet.batch_size
        return years_data, prices_data


def test():
    years = np.array([i for i in range(2000, 2022)])

    max_year = np.max(years)
    min_year = np.min(years)
    tmp_years = (years - min_year) / (max_year - min_year)

    prices = np.array(
        [10000, 11000, 12000, 13000, 14000, 12000, 13000, 16000, 18000, 20000, 19000, 22000, 24000, 23000, 26000, 35000,
         30000, 40000, 45000, 52000, 50000, 60000])
    max_price = np.max(prices)
    min_price = np.min(prices)
    tmp_prices = (prices - min_price) / (max_price - min_price)

    k = 1
    b = 0
    lr = 0.01

    batch_size = 2
    shuffle = True
    epoch = 10000

    dataset = MyDataSet(tmp_years, tmp_prices, batch_size, shuffle)

    for e in range(epoch):
        for year, price in dataset:  # 把一个对象放在for上时, 会自动调用这个对象的__iter__,
            pre = k * year + b
            loss = (pre - price) ** 2

            delta_k = 2 * (k * year + b - price) * year
            delta_b = 2 * (k * year + b - price)

            k = k - np.sum(delta_k) / batch_size * lr
            b = b - np.sum(delta_b) / batch_size * lr

    print(f"k={k}, b={b}")
    while True:
        year = (float(input("请输入年份: ")) - min_year) / (max_year - min_year)
        print("预测房价: ", (k * year + b) * (max_price - min_price) + min_price)


test()
